# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import warnings
import tabulate

from coin.strategy.accounting.fee import get_fee_rate
from coin.strategy.mm.fastfeature.ridge_regress import get_quantile
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import Ridge

# internal
from coin.support.latency.latency_model import (LatencyModel, DATA_ROOT)

# std
import datetime
from absl import app, flags

import os
import matplotlib as mpl
mpl.use("Agg")
import matplotlib.pylab as pl
import matplotlib.pyplot as plt

from coin.pnl.sim_stat_plot import setup_plt
from matplotlib.backends.backend_pdf import PdfPages


class TrModel(LatencyModel):
  def read_one_day(self, d, drop_cols):
    dstr = datetime.datetime.strftime(d, '%Y%m%d')
    path = f'{DATA_ROOT}/dataframes/taker_ratio_{flags.FLAGS.prefix}{dstr}.h5'
    if os.path.exists(path):
      print('reading file:', path)
      df = pd.read_hdf(path, 'df')
      df['filled'] = df['fill_qty'] > 0
      df['api_version'] = df['api_version'].fillna('None').astype(str)
      df['date'] = dstr
      return df
    else:
      return None

  def train_for_df(self, df, alpha=0.5, cols=None):
    assert isinstance(df, pd.DataFrame)

    y = df['filled']
    x = df[cols]

    self.uniqkeys = x[cols].drop_duplicates().reset_index(drop=True)

    self.cols = x.columns
    x = pd.get_dummies(x)
    self.dummy_cols = x.columns
    self.model = Ridge(alpha=alpha, normalize=True)
    self.model.fit(x, y.astype(float))
    self.df = df


def gen_mode(horizs):
  yield ['0ms'], ["spread_bps"], 10, "maker__spd_return"
  yield ['0ms'], ["spread_bps"], 10, "taker__spd_return"
  yield ['0ms'], ["mid_ret_bps_m3000ms_fwd"], 10, "maker__past_return"
  yield ['0ms'], ["mid_ret_bps_m3000ms_fwd"], 10, "taker__past_return"
  yield ['0ms'], ["fill_qty_ratio"], 10, "fill_qty_ratio"
  yield ['0ms'], ["fill_qty"], 10, "fill_qty"
  yield ['0ms'], ["taker_qty_ratio"], 10, "taker_qty_ratio"
  yield ['0ms'], ["symbol"], 'unique', "all"
  yield ['0ms'], ["sign"], 'unique', "sign"
  yield ['0ms'], ["model_id"], 'unique', "model_id"
  # payup: for latency matching
  # yield horizs[:6], [f"order_ret_bps_{horiz}_fwd" for horiz in horizs[:6]], 5, "payup"
  yield ['0ms'], ["lifetime_ms"], 5, "lifetime_ms"
  # mark_return: for overall performance
  yield ["subfillret"], ["subfillret"], 6, "mark_return"
  yield ['0ms'], ["signal"], 5, "signal"
  # fill qty: for lot size determination


def plot_feature_importance(df):
  pdf_filename2 = f"pta/{flags.FLAGS.prefix}_return_feature_importance_{flags.FLAGS.date}.pdf"
  pdf_filename = f"pta/{flags.FLAGS.prefix}_return_feature_{flags.FLAGS.date}.pdf"
  print(pdf_filename2)
  print(pdf_filename)
  plt.rcParams['figure.figsize'] = 6, 4
  plt.rcParams['lines.linewidth'] = 0.5
  plt.rcParams['font.size'] = 6
  plt.rcParams['legend.fontsize'] = 6
  plt.rcParams['xtick.labelsize'] = 6
  plt.rcParams['ytick.labelsize'] = 6
  columns = []
  for column in df.columns:
    if not pd.api.types.is_numeric_dtype(df[column].dtype):
      continue
    if column.find("ms_fwd") >= 0 and "order_ret_bps_m" not in column:
      continue
    if column in ["timestamp", "sign", "model_id", "filled", "maker_fee", "taker_fee"]:
      continue
    columns.append(column)
  with PdfPages(pdf_filename2) as pdf2:
    with PdfPages(pdf_filename) as pdf:
      for symbol, symboldf in df.groupby(['symbol']):
        takerorderdf = symboldf.loc[symboldf['order_ret_bps_0ms_fwd'] >= 0].reset_index(drop=True)
        ycolumn2 = '30s_midretbps'
        takerorderdf[ycolumn2] = -takerorderdf['mid_ret_bps_30000ms_fwd']
        filldf = symboldf.loc[symboldf['fill_qty'] > 0].copy()
        ycolumn = '30s_netretbps'
        filldf[ycolumn] = -filldf['fill_ret_bps_30000ms_fwd'] - filldf['net_fee_rate']
        model = RandomForestClassifier(n_estimators=10)
        X = filldf[columns]
        y = filldf[ycolumn]
        idx = ~(X.isna().any(1))
        X = X.loc[idx, :]
        y = y.loc[idx]
        model.fit(X, y > 0)
        pd.Series(model.feature_importances_, index=X.columns).nlargest(10)[::-1].plot(kind='barh')
        plt.title(symbol)
        plt.tight_layout()
        pdf2.savefig()
        plt.close()

        for column in columns:
          idx = ~np.isnan(filldf[column])
          xqi, yqi = get_quantile(filldf.loc[idx, column], filldf.loc[idx, ycolumn], 30)
          plt.plot(xqi, yqi, 'r.-', lw=0.5, markersize=0.5, rasterized=True)
          xq2, yq2 = get_quantile(takerorderdf[column], takerorderdf[ycolumn2], 30)
          plt.plot(xq2, yq2, 'g.-', lw=0.5, markersize=0.5, rasterized=True)
          plt.legend([ycolumn, ycolumn2])
          plt.xlabel(column)
          plt.ylabel(ycolumn)
          plt.title(symbol)
          plt.grid()
          plt.tight_layout()
          pdf.savefig(dpi=100)
          plt.close()


def main(_):
  os.makedirs("pta", exist_ok=True)

  warnings.simplefilter(action='ignore', category=FutureWarning)
  warnings.simplefilter(action='ignore', category=UserWarning)

  model = TrModel()

  columnstr = 'exchange,market_type,symbol,date'
  columns = columnstr.split(",")
  model.train_for_dates(flags.FLAGS.date, columnstr, ycol='filled')
  print(model.dummy_cols)

  df = model.df
  df['fill_qty_ratio'] = df.fill_qty / df.order_qty
  df['taker_qty_ratio'] = df.taker_fill_qty / df.fill_qty

  for market_type, exchange, symbol, date in model.df.groupby(
      ['market_type', 'exchange', 'symbol', 'date']).first().index:
    start_time = pd.DatetimeIndex([date])[0].to_pydatetime()
    product = generate_product_from_str2(market_type=market_type,
                                         exchange=exchange,
                                         api_version=None,
                                         norm_product=symbol,
                                         current_datetime=start_time)
    maker_fee, taker_fee = get_fee_rate(market_type, exchange, product)
    selidx = model.df[['market_type', 'exchange', 'symbol']] == (market_type, exchange, symbol)
    chosen = selidx.all(1)
    model.df.loc[chosen, 'maker_fee'] = maker_fee
    model.df.loc[chosen, 'taker_fee'] = taker_fee

  rows = []

  horizs = [
      'm10000ms',
      'm3000ms',
      'm1000ms',
      '0ms',
      '25ms',
      '50ms',
      '100ms',
      '200ms',
      '500ms',
      '1000ms',
      '3000ms',
      '5000ms',
      '10000ms',
      '30000ms',
      '60000ms'
  ]

  rainbow_cm = pl.cm.rainbow
  winter_cm = pl.cm.winter

  columnstr = 'exchange,market_type,symbol'
  columns = columnstr.split(",")
  df['net_fee_rate'] = (
      df['taker_fill_qty'] * df['taker_fee'] + (
          df['fill_qty'] - df['taker_fill_qty']) * df['maker_fee']) / df['fill_qty']
  plot_feature_importance(df)

  prtdfs = []
  for markrethorizs, markretnames, nqtl, filename in gen_mode(horizs):
    if filename == 'signal' and np.isnan(df.signal).all():
      print("nan signals(maybe didn't dump).. omitting.")
      continue
    pdf_filename = f"pta/{flags.FLAGS.prefix}{filename}_{flags.FLAGS.date}.pdf"
    print(pdf_filename)
    with PdfPages(pdf_filename) as pdf:
      if filename.startswith('maker__'):
        target_df = df.loc[(df.fill_qty > 0) & (df.taker_fill_qty == 0)]
      elif filename.startswith('taker__'):
        target_df = df.loc[(df.taker_fill_qty > 0) & ((df.taker_fill_qty - df.fill_qty) == 0)]
      elif filename == 'taker_qty_ratio':
        target_df = df.loc[df.fill_qty > 0]
      else:
        target_df = df

      for key, subdf in target_df.groupby(columns):
        date_range = subdf['date'].min() + '-' + subdf['date'].max()
        if len(subdf) < 10:
          print(f"{key}: # too small, omitted: {len(subdf)}")
          continue
        for markrethoriz, markretname in zip(markrethorizs, markretnames):
          if nqtl == 'unique':
            ncount = len(subdf[markretname].unique())
            colors = rainbow_cm(np.linspace(0, 1, ncount))
          else:
            colors = winter_cm(np.linspace(0, 1, nqtl))

          subdf = subdf.reset_index(drop=True)
          qtls = []
          subrows = collections.OrderedDict()
          if markrethoriz in ["subfillret"]:
            for horiz in horizs:
              colname = f'order_ret_bps_{horiz}_fwd'
              colname2 = f'fill_ret_bps_{horiz}_fwd'
              colname3 = f'order_payup_bps_{horiz}_fwd'
              colname4 = f'payup0_ret_bps_{horiz}_fwd'
              sub_ret_mean = subdf[colname].mean()
              sub_bbo_ret_mean = subdf[colname3].mean()
              payup_ret_mean = subdf[colname4].mean()
              fee_total = subdf['taker_fill_qty'] * subdf['taker_fee'] + (
                  subdf['fill_qty'] - subdf['taker_fill_qty']) * subdf['maker_fee']
              fee_mean = fee_total.sum() / subdf['fill_qty'].sum()
              if subdf['filled'].sum() == 0:
                fill_ret_mean = 0
                fill_act_ret_mean = 0
              else:
                fill_ret_mean = np.average(subdf[colname], weights=subdf['filled'])
                fill_act_ret_mean = np.average(subdf[colname2][subdf['filled']],
                                               weights=subdf['fill_qty'][subdf['filled']])
              qtls.append({
                  'i': 0,
                  'dt': horiz,
                  'col': 'orderp[t]-bbo[t+dt]',
                  'orderp[t]-bbo[t+dt]': sub_bbo_ret_mean
              })
              qtls.append({
                  'i': 1,
                  'dt': horiz,
                  'col': 'bbo[t]-mid[t+dt]',
                  'bbo[t]-mid[t+dt]': payup_ret_mean
              })
              qtls.append({
                  'i': 2,
                  'dt': horiz,
                  'col': 'orderp[t]-mid[t+dt]',
                  'orderp[t]-mid[t+dt]': sub_ret_mean
              })
              qtls.append({
                  'i': 3,
                  'dt': horiz,
                  'col': 'orderp[t]-mid[t+dt]',
                  'orderp[t]-mid[t+dt]': fill_ret_mean
              })
              qtls.append({
                  'i': 4,
                  'dt': horiz,
                  'col': 'fillp[t]-mid[t+dt]',
                  'fillp[t]-mid[t+dt]': fill_act_ret_mean
              })
              qtls.append({'i': 5, 'dt': horiz, 'col': 'fee_mean', 'fee_mean': -fee_mean})
          else:
            if nqtl == 'unique':
              if filename == 'sign':
                # simply [-1, 1] ordered
                quantiles = subdf[markretname].astype('category')
              else:
                quantiles = subdf[markretname].astype('category').cat.as_ordered()
                nsamples = subdf[markretname].groupby(quantiles).count()
                quantiles = quantiles.cat.reorder_categories(nsamples.sort_values().index.tolist(),
                                                             ordered=True)
              nsamples = subdf[markretname].groupby(quantiles).count()
              quantiles_val = subdf[markretname].groupby(quantiles).first()
            else:
              subdf['rank'] = subdf[markretname].rank(method='first')
              quantiles = pd.qcut(subdf['rank'], nqtl)
              nsamples = subdf[markretname].groupby(quantiles).count()
              quantiles_val = subdf[markretname].groupby(quantiles).mean()
            frm_means = ((subdf['fill_qty'] > 0).astype(float).groupby(quantiles).mean() /
                         (subdf['order_qty'] > 0).astype(float).groupby(quantiles).mean())
            frq_means = (subdf['fill_qty'].groupby(quantiles).mean()
                         / subdf['order_qty'].groupby(quantiles).mean())
            tfr_means = (subdf['taker_fill_qty'].groupby(quantiles).mean()
                         / subdf['fill_qty'].groupby(quantiles).mean())
            fee_paid = subdf['taker_fill_qty'] * subdf['taker_fee'] + (
                subdf['fill_qty'] - subdf['taker_fill_qty']) * subdf['maker_fee']
            fee_means = (
                fee_paid * 1e4
                    ).groupby(quantiles).mean() / subdf['fill_qty'].groupby(quantiles).mean()
            for horiz in horizs:
              colname = f'order_ret_bps_{horiz}_fwd'
              colname2 = f'fill_ret_bps_{horiz}_fwd'
              ret_means = subdf[colname].groupby(quantiles).mean()
              fill_act_ret_means = []
              fill_act_pnls = []
              for _, groupdf in subdf.groupby(quantiles):
                if groupdf['filled'].sum() == 0:
                  fill_act_ret_mean = 0
                  fill_act_pnl = 0
                else:
                  fill_act_ret_mean = np.average(groupdf[colname2][groupdf['filled']],
                                                 weights=groupdf['fill_qty'][groupdf['filled']])
                  feeseri = groupdf['taker_fill_qty'] * groupdf['taker_fee'] + (
                      groupdf['fill_qty'] - groupdf['taker_fill_qty']) * groupdf['maker_fee']
                  fill_act_pnl = np.sum(
                      -groupdf[colname2][groupdf['filled']]
                      * groupdf['fill_qty'][groupdf['filled']]) - feeseri.sum() * 1e4
                fill_act_ret_means.append(fill_act_ret_mean)
                fill_act_pnls.append(fill_act_pnl)
              for i, (qtl_mean, frq_mean, frm_mean, tfr_mean, ret_mean, fill_act_ret_mean,
                      fill_act_pnl, fee_mean, nsample) in enumerate(
                          zip(quantiles_val, frq_means, frm_means, tfr_means, ret_means,
                              fill_act_ret_means, fill_act_pnls, fee_means, nsamples)):
                qtls.append({
                    'i': i,
                    'dt': horiz,
                    'col': 'orderp[t]-mid[t+dt]',
                    'qtlcol': markretname,
                    markretname: qtl_mean,
                    'orderp[t]-mid[t+dt]': ret_mean,
                    'fillp[t]-mid[t+dt]': fill_act_ret_mean,
                    'fill_act_pnl': fill_act_pnl,
                    'netret': -fill_act_ret_mean - fee_mean,
                    'frq': frq_mean,
                    'frm': frm_mean,
                    'tfr': tfr_mean,
                    'feemean': fee_mean,
                    'nsample': nsample
                })
          setup_plt()
          if markrethoriz is not None and markrethoriz not in ["subfillret"]:
            plt.axvline(x=horizs.index(markrethoriz))
            qtldf = pd.DataFrame(qtls)

            if 'col' not in qtldf:
              continue

            retcols = qtldf['col'].unique()
            assert len(retcols) == 1
            retcol = retcols[0]
            retsel = qtldf['col'] == retcol
            netret60sqtl = qtldf[(qtldf.dt == '60000ms') & retsel]['netret']
            netret30sqtl = qtldf[(qtldf.dt == '30000ms') & retsel]['netret']
            netret5sqtl = qtldf[(qtldf.dt == '5000ms') & retsel]['netret']
            ordret0sqtl = -qtldf[(qtldf.dt == '0ms') & retsel]['orderp[t]-mid[t+dt]']
            ordret5sqtl = -qtldf[(qtldf.dt == '5000ms') & retsel]['orderp[t]-mid[t+dt]']
            qtldfsel = qtldf[retsel]

            nsamples = ["%6.0f" % val for val in qtldfsel.groupby('i').first()['nsample']]

            net60spnlqtl = qtldf[(qtldf.dt == '60000ms') & retsel]['fill_act_pnl']
            net60spnls = ["%6.0f" % net60spnlq for net60spnlq in net60spnlqtl]

            qtlcol = qtldf['qtlcol'].unique()[0]
            valqtl = qtldf[(qtldf.dt == markrethoriz) & retsel][qtlcol]
            if nqtl == 'unique':
              qtl_means = ["%s" % val for val in valqtl]
            else:
              qtl_means = ["%6.2f" % val for val in valqtl]
            orderret_0s = ["%6.2f" % val for val in ordret0sqtl]
            orderret_5s = ["%6.2f" % val for val in ordret5sqtl]
            net_ret5s = ["%6.2f" % val for val in netret5sqtl]
            net_ret30s = ["%6.2f" % val for val in netret30sqtl]
            net_ret60s = ["%6.2f" % val for val in netret60sqtl]
            frq_means = ["%4.2f" % val for val in qtldfsel.groupby('i').first()['frq']]
            frm_means = ["%4.2f" % val for val in qtldfsel.groupby('i').first()['frm']]
            taker_ratios = ["%6.2f" % val for val in qtldfsel.groupby('i').first()['tfr']]
            fee_means = ["%6.2f" % val for val in qtldfsel.groupby('i').first()['feemean']]

          for i, qtldf in pd.DataFrame(qtls).groupby('i'):
            qtldf = qtldf.reset_index(drop=True)
            retcol = qtldf.col.unique()
            assert len(retcol) == 1
            retcol = retcol[0]
            plt.plot(qtldf['dt'],
                     -qtldf[retcol],
                     linestyle='dashed',
                     marker='.',
                     lw=0.5,
                     markersize=1,
                     color=colors[i])
            if 'fr' in qtldf:
              fr = qtldf['fr']
              plt.scatter(qtldf['dt'], -qtldf[retcol], c=(1 - fr), s=10, alpha=0.5)
            plt.gray()
          if markrethoriz in ["subfillret"]:
            plt.legend([
                "subsret(orderp to bbop)",
                "subsret(bbop to midp = payup0 ret)",
                "subsret(orderp to midp)",
                "fillret(orderp to midp)",
                "fillret(fill to midp)",
                "fee_mean",
            ], loc='upper left')
          else:
            adf = pd.DataFrame([
                qtl_means,
                orderret_0s,
                orderret_5s,
                net_ret5s,
                net_ret30s,
                net_ret60s,
                net60spnls,
                frm_means,
                frq_means,
                taker_ratios,
                fee_means,
                nsamples],
                columns=range(len(qtl_means)),
                index=[
                    qtlcol,
                    'ordp0smid',
                    'ordp5smid',
                    'net5smid',
                    'net30smid',
                    'net60smid',
                    'net60spnl',
                    'FRmsg',
                    'FRqty',
                    'TakR',
                    'fee',
                    '#sample'])
            plt.legend(adf.T.to_string().split("\n"), loc='upper left', fontsize=5)
            prtdf = adf.T[['net5smid', 'net30smid', 'net60smid', 'net60spnl']]
            prtdf['key'] = ":".join(key)
            prtdf['col'] = qtlcol
            prtdf['qtl'] = adf.T[qtlcol]
            prtdf = prtdf[['col', 'qtl', 'key', 'net5smid', 'net30smid', 'net60smid', 'net60spnl']]
            prtdfs.append(prtdf)
          plt.ylabel('(return to orderp): mid[t+dt]-orderp[t]')
          plt.grid(lw=0.5, which='major')
          if markrethoriz is not None:
            plt.title(f"mark.ret {markrethoriz} {key} {date_range}", fontsize=8)
        pdf.savefig()
        plt.close()
        rows.extend(subrows.values())
  prt = "\n".join([tabulate.tabulate(prtdf, headers='keys', tablefmt="psql") for prtdf in prtdfs])
  wfile = open(f"pta/{flags.FLAGS.prefix}_{flags.FLAGS.date}.txt", "w")
  print(prt, file=wfile)


if __name__ == "__main__":
  flags.DEFINE_string('date', None, 'example: --date 20190708')
  flags.DEFINE_string('prefix', '', '')
  flags.mark_flag_as_required('date')
  app.run(main)
